package initialize

import (
	"colly/pkg/hook"
	"context"
	"fmt"
	"github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
	"io"
	"os"
	"time"
)

func NewGorm(dsn, prefix, mode string) *Gorm {
	return &Gorm{
		Dsn:    dsn,
		Prefix: prefix,
		Mode:   mode,
	}
}

func getLogger(ctx context.Context) *logrus.Logger {
	fileLogger := &lumberjack.Logger{
		Filename:   "./logs/sql.log",
		MaxSize:    10,   // Max size in MB
		MaxBackups: 10,   // Max number of old log files to keep
		MaxAge:     30,   // Max age in days to keep old log files
		Compress:   true, // Compress old log files
	}
	lg := logrus.New()
	lg.SetOutput(io.MultiWriter(os.Stdout, fileLogger))
	lg.SetFormatter(&logrus.JSONFormatter{
		DisableTimestamp: false,
		TimestampFormat:  "2006-01-02 15:04:05",
		FieldMap:         nil,
		CallerPrettyfier: nil,
	})
	lg.AddHook(&hook.ContextHook{Context: ctx})
	return lg
}

type CustomGORMLogger struct {
	logrusLogger *logrus.Logger
	config       logger.Config
}

func NewCustomGORMLogger(logrusLogger *logrus.Logger, config logger.Config) *CustomGORMLogger {
	return &CustomGORMLogger{
		logrusLogger: logrusLogger,
		config:       config,
	}
}

func (c *CustomGORMLogger) LogMode(level logger.LogLevel) logger.Interface {
	newLogger := *c
	switch level {
	case logger.Silent:
		newLogger.logrusLogger.SetLevel(logrus.PanicLevel)
	case logger.Error:
		newLogger.logrusLogger.SetLevel(logrus.ErrorLevel)
	case logger.Warn:
		newLogger.logrusLogger.SetLevel(logrus.WarnLevel)
	case logger.Info:
		newLogger.logrusLogger.SetLevel(logrus.InfoLevel)
	default:
		newLogger.logrusLogger.SetLevel(logrus.InfoLevel)
	}
	return &newLogger
}

func (c *CustomGORMLogger) Info(ctx context.Context, msg string, data ...interface{}) {
	c.logrusLogger.WithContext(ctx).Infof(msg, data...)
}

func (c *CustomGORMLogger) Warn(ctx context.Context, msg string, data ...interface{}) {
	c.logrusLogger.WithContext(ctx).Warnf(msg, data...)
}

func (c *CustomGORMLogger) Error(ctx context.Context, msg string, data ...interface{}) {
	c.logrusLogger.WithContext(ctx).Errorf(msg, data...)
}

func (c *CustomGORMLogger) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
	elapsed := time.Since(begin)
	switch {
	case err != nil && c.config.LogLevel >= logger.Error:
		sql, rows := fc()
		if rows == -1 {
			c.logrusLogger.WithContext(ctx).Errorf("%s [%.3fms] %s", err, float64(elapsed.Nanoseconds())/1e6, sql)
		} else {
			c.logrusLogger.WithContext(ctx).Errorf("%s [%.3fms] [rows:%v] %s", err, float64(elapsed.Nanoseconds())/1e6, rows, sql)
		}
	case elapsed > c.config.SlowThreshold && c.config.SlowThreshold != 0 && c.config.LogLevel >= logger.Warn:
		sql, rows := fc()
		slowLog := fmt.Sprintf("SLOW SQL >= %v", c.config.SlowThreshold)
		if rows == -1 {
			c.logrusLogger.WithContext(ctx).Warnf("%s [%.3fms] %s", slowLog, float64(elapsed.Nanoseconds())/1e6, sql)
		} else {
			c.logrusLogger.WithContext(ctx).Warnf("%s [%.3fms] [rows:%v] %s", slowLog, float64(elapsed.Nanoseconds())/1e6, rows, sql)
		}
	case c.config.LogLevel == logger.Info:
		sql, rows := fc()
		if rows == -1 {
			c.logrusLogger.WithContext(ctx).Infof("[%.3fms] %s", float64(elapsed.Nanoseconds())/1e6, sql)
		} else {
			c.logrusLogger.WithContext(ctx).Infof("[%.3fms] [rows:%v] %s", float64(elapsed.Nanoseconds())/1e6, rows, sql)
		}
	}
}

type Gorm struct {
	Dsn     string
	Prefix  string
	Mode    string
	Context context.Context
}

func (r Gorm) Connect(ctx context.Context) *gorm.DB {
	//customLogger := NewCustomGORMLogger(getLogger(ctx), logger.Config{
	//	SlowThreshold:             time.Second,
	//	LogLevel:                  logger.Info,
	//	IgnoreRecordNotFoundError: true,
	//	Colorful:                  false,
	//})

	db, err := gorm.Open(mysql.New(mysql.Config{
		DSN:                       r.Dsn, // DSN data source name
		DefaultStringSize:         256,   // string 类型字段的默认长度
		DisableDatetimePrecision:  true,  // 禁用 datetime 精度，MySQL 5.6 之前的数据库不支持
		DontSupportRenameIndex:    true,  // 重命名索引时采用删除并新建的方式，MySQL 5.7 之前的数据库和 MariaDB 不支持重命名索引
		DontSupportRenameColumn:   true,  // 用 `change` 重命名列，MySQL 8 之前的数据库和 MariaDB 不支持重命名列
		SkipInitializeWithVersion: false, // 根据当前 MySQL 版本自动配置
	}), &gorm.Config{
		PrepareStmt: true,
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   r.Prefix, // 表明前缀，可不设置
			SingularTable: true,     // 使用单数表名，即不会在表名后添加复数s
		},
		Logger: logger.Default.LogMode(logger.Error),
	})
	if err != nil {
		panic(err)
	}
	return db
}
